import base64
import hashlib

from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_v1_5 as PKCS1_cipher, AES


class RSAEncryptor:
    @classmethod
    def get_key(cls, key_file):
        with open(key_file) as f:
            data = f.read()  # 获取，密钥信息
            key = RSA.importKey(data)
        return key

    @classmethod
    def encrypt_data(cls, public_key, msg):
        public_key = cls.get_key(public_key)  # 读取公钥信息
        cipher = PKCS1_cipher.new(public_key)  # 生成一个加密的类
        encrypt_text = base64.b64encode(cipher.encrypt(msg.encode()))  # 对数据进行加密
        return encrypt_text.decode()  # 对文本进行解码码

    @classmethod
    def decrypt_data(cls, public_key, encrypt_msg):
        private_key = cls.get_key(public_key)  # 读取私钥信息
        cipher = PKCS1_cipher.new(private_key)  # 生成一个解密的类
        back_text = cipher.decrypt(base64.b64decode(encrypt_msg), 0)  # 进行解密
        return back_text.decode()


class AESEncryptor:

    @classmethod
    def pkcs7padding(cls, text):
        """
        明文使用PKCS7填充
        最终调用AES加密方法时，传入的是一个byte数组，要求是16的整数倍，因此需要对明文进行处理
        :param text: 待加密内容(明文)
        :return:
        """
        bs = AES.block_size  # 16
        length = len(text)
        bytes_length = len(bytes(text, encoding='utf-8'))
        # tips：utf-8编码时，英文占1个byte，而中文占3个byte
        padding_size = length if (bytes_length == length) else bytes_length
        padding = bs - padding_size % bs
        # tips：chr(padding)看与其它语言的约定，有的会使用'\0'
        padding_text = chr(padding) * padding
        return text + padding_text

    @classmethod
    def encode(cls, text, key):
        aes = AES.new(str.encode(key), AES.MODE_ECB)  # 初始化加密器
        content_padding = AESEncryptor.pkcs7padding(text)
        # 加密
        encrypt_bytes = aes.encrypt(bytes(content_padding, encoding='utf-8'))
        # 重新编码
        result = str(base64.b64encode(encrypt_bytes), encoding='utf-8')
        return result

    @staticmethod
    def decode(text, key):
        aes = AES.new(str.encode(key), AES.MODE_ECB)  # 初始化加密器
        decrypted_text = aes.decrypt(base64.decodebytes(bytes(text, encoding='utf8'))).decode("utf8")  # 解密
        decrypted_text = decrypted_text[:-ord(decrypted_text[-1])]  # 去除多余补位
        return decrypted_text


def sha256_hash(text, salt):
    text = f'{text}{salt}'
    return hashlib.sha256(text.encode()).hexdigest()


if __name__ == '__main__':
    print(RSAEncryptor.encrypt_data('../../auth-public.pem', '1111'))

